import os
import cv2
import numpy as np
#2020.3.23's calibration
from volumeMeasure import stereo_calibrator

class Calib1:
    def __init__(self):
        self.pic_dir = 'C:/Users/tellw/Desktop/bishe/items/calibed_pics/20200323115211'
        self.two_frame_w=1280
        self.two_frame_h=480
        self.pic_pattern=(6, 4)
        self.left_intrinsic = np.eye(3, dtype=np.float64)
        self.right_intrinsic = np.eye(3, dtype=np.float64)
        self.left_coeffs = np.zeros((1, 5), dtype=np.float64)
        self.right_coeffs = np.zeros((1, 5), dtype=np.float64)
        self.rotation = np.zeros((3, 3), dtype=np.float64)
        self.transform = np.zeros((3, 3), dtype=np.float64)
        self.E = np.zeros((3, 3), dtype=np.float64)
        self.F = np.zeros((3, 3), dtype=np.float64)
        self.r1 = np.zeros((3, 3), dtype=np.float64)
        self.r2 = np.zeros((3, 3), dtype=np.float64)
        self.p1 = np.zeros((3, 4), dtype=np.float64)
        self.p2 = np.zeros((3, 4), dtype=np.float64)
        self.q = np.zeros((4, 4), dtype=np.float64)
        self.bStereoCalib = True
        self.image_size = (self.two_frame_h, self.two_frame_w//2)

    def test_calib(self, bShow=True, bRecordedToFile=False):
        world_coordis = self._cal_corner_positions((6, 4), (39.9, 39.8))
        # print(world_coordis)
        # print(world_coordis.shape, world_coordis.dtype)
        # input()
        left_corners = []
        right_corners = []
        world_coordiss = [world_coordis for i in range(54)]
        #print(world_coordiss)
        for pic in os.listdir(self.pic_dir):
            if pic.endswith('.jpg'):
                two_frame = cv2.imread(os.path.join(self.pic_dir, pic))
                left = two_frame[:, :(self.two_frame_w//2), :]
                right = two_frame[:, (self.two_frame_w//2):, :]
                retl, left_corner = cv2.findChessboardCorners(left, self.pic_pattern, None)
                cv2.cornerSubPix(cv2.cvtColor(left, cv2.COLOR_BGR2GRAY), left_corner, (5, 5), (-1, -1), (cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_MAX_ITER, 30, 0.1))
                left_corners.append(left_corner)
                if retl and bRecordedToFile:
                    with open(os.path.join(self.pic_dir, 'left_pic_points.txt'), 'a') as f:
                        f.write(str(left_corner)+'\n')
                retr, right_corner = cv2.findChessboardCorners(right, self.pic_pattern, None)
                cv2.cornerSubPix(cv2.cvtColor(right, cv2.COLOR_BGR2GRAY), right_corner, (5, 5), (-1, -1), (cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_MAX_ITER, 30, 0.1))
                right_corners.append(right_corner)
                if retr and bRecordedToFile:
                    with open(os.path.join(self.pic_dir, 'right_pic_points.txt'), 'a') as f:
                        f.write(str(right_corner)+'\n')
                if bShow:
                    cv2.drawChessboardCorners(left, self.pic_pattern, left_corner, retl)
                    cv2.drawChessboardCorners(right, self.pic_pattern, right_corner, retr)
                    cv2.imshow('left_corners_show', left)
                    cv2.imshow('right_corners_show', right)
                    key = cv2.waitKey()
                    if key == ord('n'):
                        os.rename(os.path.join(self.pic_dir, pic), os.path.join(self.pic_dir, 'not_ok', pic))
                    elif key == 27 or key == ord('q'):
                        break
        calib_left_ret, self.left_intrinsic, self.left_coeffs, left_rvecs, left_tvecs = cv2.calibrateCamera(
            world_coordiss, left_corners, self.image_size, self.left_intrinsic, self.left_coeffs, flags=cv2.CALIB_FIX_K3)
        calib_right_ret, self.right_intrinsic, self.right_coeffs, right_rvecs, right_tvecs = cv2.calibrateCamera(
            world_coordiss, right_corners, self.image_size, self.right_intrinsic, self.right_coeffs, flags=cv2.CALIB_FIX_K3)
        print('left calibration error: {}, right calibration error: {}'.format(calib_left_ret, calib_right_ret))
        if self.bStereoCalib:
            stereo_calib_ret, self.left_intrinsic, self.left_coeffs, self.right_intrinsic, self.right_coeffs, self.rotation, self.transform, self.E, self.F = cv2.stereoCalibrate(
                world_coordiss, left_corners, right_corners, self.left_intrinsic, self.left_coeffs,
                self.right_intrinsic, self.right_coeffs, self.image_size)
            print('stereo calibration error: {}'.format(stereo_calib_ret))
            self.r1, self.r2, self.p1, self.p2, self.q, self.roi1, self.roi2 = cv2.stereoRectify(cameraMatrix1=self.left_intrinsic, distCoeffs1=self.left_coeffs, cameraMatrix2=self.right_intrinsic, distCoeffs2=self.right_coeffs, imageSize=self.image_size, R=self.rotation, T=self.transform)
        self.save_params('newStereoCalib.xml')

    def save_params(self, xmlfile):
        storage = cv2.FileStorage(xmlfile, cv2.FILE_STORAGE_WRITE + cv2.FILE_STORAGE_FORMAT_XML)
        if not storage.isOpened():
            print('storage is not opened')
            return
        storage.write('left_intrinsic', self.left_intrinsic)
        storage.write('left_coeffs', self.left_coeffs)
        storage.write('right_intrinsic', self.right_intrinsic)
        storage.write('right_coeffs', self.right_coeffs)
        if self.bStereoCalib:
            storage.write('rotation', self.rotation)
            storage.write('transform', self.transform)
            storage.write('E', self.E)
            storage.write('F', self.F)
            storage.write('r1', self.r1)
            storage.write('r2', self.r2)
            storage.write('p1', self.p1)
            storage.write('p2', self.p2)
            storage.write('q', self.q)
            storage.write('roi1', self.roi1)
            storage.write('roi2', self.roi2)
        storage.release()

    def _cal_corner_positions(self, pattern, grid_size):
        result = []
        for row in range(pattern[1]):
            for col in range(pattern[0]):
                result.append([col*grid_size[0], row*grid_size[1], 0])#x轴，y轴分别是图像坐标系的两个轴，z轴垂直于地面竖直向上，grid_size还是按照图像坐标系的习惯索引0即是x轴方向的单网格长度，索引1即是y轴方向的单网格长度
        return np.array(result).astype(np.float32)

if __name__ == '__main__':
    # c = Calib1()
    # c.test_calib(False)
    xmlfile = 'parametersOfCamera/selfMakeParam1910262022.xml'
    s = stereo_calibrator(xmlfile, 'static')
    r1, r2, p1, p2, q, roi1, roi2 = cv2.stereoRectify(
        s.left_intrinsic, s.left_coeffs, s.right_intrinsic,
        s.right_coeffs, (640, 480), s.rotation, s.transform)
    storage = cv2.FileStorage('lala.xml', cv2.FILE_STORAGE_WRITE + cv2.FILE_STORAGE_FORMAT_XML)
    if not storage.isOpened():
        print('storage is not opened')
    storage.write('r1', r1)
    storage.write('r2', r2)
    storage.write('p1', p1)
    storage.write('p2', p2)
    storage.write('q', q)
    storage.write('roi1', roi1)
    storage.write('roi2', roi2)
    storage.release()